package com.dianping.tpfun.supervised.decisiontree.algorithm;

import com.dianping.tpfun.supervised.decisiontree.bean.*;
import com.google.common.collect.Lists;
import com.google.common.collect.Maps;
import org.apache.commons.collections.CollectionUtils;


import java.util.Collections;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;

/**
 * C4.5实现
 * Created by ztw on 16/1/6.
 */
public class DecisionTreeC {

    public static double THRESHOLD = 0.01;

    public DefaultDecisionTreeNode train(List<DefaultData> datas, List<DefaultFeature> features) {

        DefaultDecisionTreeNode currentNode = new DefaultDecisionTreeNode();
        if (datas == null || datas.isEmpty()) {
            return null;
        }

        if (checkDataLabelsUnique(datas)) {
            currentNode.setIsLeaf(true);
            currentNode.setLabel(datas.get(0).getLabel());
            return currentNode;
        }

        if (features.isEmpty()) {
            currentNode.setIsLeaf(true);
            //取多数类
            currentNode.setLabel(datas.get(0).getLabel());
            return currentNode;
        }

        //选择最优特征
        DefaultFeature chooseFeature = chooseFeatureAndPartion(datas, features);
        Map<Comparable, List<DefaultData>> partionsMap = chooseFeature.getPartionDatas();
        currentNode.setPartionPoints(Lists.newLinkedList(partionsMap.keySet()));
        LinkedHashMap<Comparable, DefaultDecisionTreeNode> partionChildMap = Maps.newLinkedHashMap();
        currentNode.setPartionChildMap(partionChildMap);
        features.remove(chooseFeature);

        for (Map.Entry<Comparable, List<DefaultData>> entry : partionsMap.entrySet()) {
            DefaultDecisionTreeNode child = train(entry.getValue(), features);
            partionChildMap.put(entry.getKey(), child);
        }
        features.add(chooseFeature);
        return currentNode;
    }

    private boolean checkDataLabelsUnique(List<DefaultData> datas) {
        DefaultLabel label = datas.get(0).getLabel();
        for (DefaultData data : datas) {
            if (!data.getLabel().equals(label)) {
                return false;
            }
        }
        return true;
    }

    public DefaultFeature chooseFeatureAndPartion(List<DefaultData> datas, List<DefaultFeature> features) {

        double maxScore = Double.MIN_VALUE;
        DefaultFeature chooseFeature = null;
        Map<Comparable, List<DefaultData>> choosePartion = null;
        for (DefaultFeature feature : features) {
            Map<Comparable, List<DefaultData>> partions = Maps.newLinkedHashMap();
            double score = calculateInfoGainRatio(datas, feature, partions);
            if (score > maxScore) {
                chooseFeature = feature;
                maxScore = score;
            }
            choosePartion = partions;
        }
        chooseFeature.setPartionDatas(choosePartion);
        return chooseFeature;
    }

    public double calculateInfoGainRatio(List<DefaultData> datas, DefaultFeature feature, Map<Comparable, List<DefaultData>> partionMap) {

        int featureIndex = feature.getIndex();
        double result = 0.0;
        //信息增益比的分母部分
        double featureEntropy = 0.0;
        //离散特征处理
        if (feature.getType() == FeatureType.Discrete.getValue()) {
            for (DefaultData defaultData : datas) {
                Comparable dataFeatureValue = defaultData.getDataDims()[featureIndex].getData();
                if (partionMap.containsKey(dataFeatureValue)) {
                    List<DefaultData> tempList = partionMap.get(dataFeatureValue);
                    tempList.add(defaultData);
                } else {
                    partionMap.put(dataFeatureValue, Lists.asList(defaultData, null));
                }
            }

        }
        else if (feature.getType() == FeatureType.Continuity.getValue()) {
            //连续特征处理
            List<Comparable> partionPoint = getPartionPoint(datas,feature);
        }
        result = caculateFeatureInfoGain(partionMap, datas, feature , caculatePreInfoGain(datas));

        return result;
    }

    /**
        该方法处理连续变量的分割以及打分问题
     */
    public <T extends Comparable> List<T> getPartionPoint(List<DefaultData> datas, DefaultFeature<T> feature){
        List<T> candidatePartionPointList = Lists.newLinkedList();
        List<T> resultPartionPointList = Lists.newLinkedList();

        if(CollectionUtils.isEmpty(datas)||datas.size()<2){
            return resultPartionPointList;
        }

        int featureIndex = feature.getIndex();
        for(int i=1;i<datas.size();i++){
            DefaultData cur = datas.get(i);
            DefaultData pre = datas.get(i-1);
            if(cur.getLabel()!= pre.getLabel()){
                Comparable curFeatureData = cur.getDataDims()[featureIndex].getData();
                Comparable preFeatureData = pre.getDataDims()[featureIndex].getData();
                candidatePartionPointList.add((T)preFeatureData);
            }
        }
        double maxGain = Double.MIN_VALUE;
        int index = 0;

        //存放左右分区数据
        List<DefaultData> leftPartions=Lists.newLinkedList();
        List<DefaultData> rightPartions=Lists.newLinkedList(datas);

        //选择最优分割点
        for(T candidatePartionPoint : candidatePartionPointList){
            Map<Comparable, List<DefaultData>> partionData = Maps.newHashMap();
            //遍历推进数据
            while (datas.get(index).getDataDims()[featureIndex].getData().compareTo(candidatePartionPoint)<=0){
                leftPartions.add(datas.get(index));
                rightPartions.remove(datas.get(index));
                index++;
            }
            partionData.put(candidatePartionPoint ,leftPartions);
            //二分的时候分割点只有一个，暂时伪造另外一个分割点
            partionData.put(new Comparable() {
                @Override
                public int compareTo(Object o) {
                    return 0;
                }
            }, leftPartions);

            double candidateInfoGain = caculateFeatureInfoGain(partionData,datas,feature,caculatePreInfoGain(datas));
            if(candidateInfoGain>maxGain){
                resultPartionPointList.clear();
                resultPartionPointList.add(candidatePartionPoint);
            }
        }
        return resultPartionPointList;
    }

    public double caculateFeatureInfoGain(Map<Comparable, List<DefaultData>> partionData, List<DefaultData> originDatas,
                                          DefaultFeature feature, double preInfoGain) {

        double result = 0.0;
        double featureValueProbabilitys[] = new double[partionData.size()];
        double conditionFeatureEntropy[] = new double[partionData.size()];
        int index = 0;
        for (Map.Entry<Comparable, List<DefaultData>> classEntry : partionData.entrySet()) {
            featureValueProbabilitys[index] = classEntry.getValue().size() / originDatas.size();
            conditionFeatureEntropy[index] = caculatePreInfoGain(classEntry.getValue());
            result += featureValueProbabilitys[index] * conditionFeatureEntropy[index];
            index++;
        }

        return preInfoGain-result;
    }

    public double caculatePreInfoGain(List<DefaultData> datas) {

        Map<DefaultLabel, List<DefaultData>> classPartionMap = Maps.newHashMap();
        for (DefaultData defaultData : datas) {
            DefaultLabel dataLabel = defaultData.getLabel();
            if (classPartionMap.containsKey(dataLabel)) {
                List<DefaultData> tempList = classPartionMap.get(dataLabel);
                tempList.add(defaultData);
            } else {
                classPartionMap.put(dataLabel, Lists.asList(defaultData, null));
            }
        }

        double classProbabilitys[] = new double[classPartionMap.size()];
        int index = 0;
        for (Map.Entry<DefaultLabel, List<DefaultData>> classEntry : classPartionMap.entrySet()) {
            classProbabilitys[index] = classEntry.getValue().size() / datas.size();
            index++;
        }
        double result = 0.0;
        for (double classProbability : classProbabilitys) {
            result += classProbability * Math.log(classProbability);
        }
        return result * (-1);
    }
}
